import subprocess
import sys
import logging

logger = logging.getLogger(__name__)

class CodeExecutionError(Exception):
    """Custom exception for code execution errors."""
    def __init__(self, message, stderr):
        super().__init__(message)
        self.stderr = stderr

def execute(code: str, input_json: str, timeout: int = 30) -> str:
    """
    Executes a string of Python code in a sandboxed subprocess.

    Args:
        code: The Python code to execute.
        input_json: A JSON string that will be passed to the subprocess's stdin.
        timeout: The maximum execution time in seconds.

    Returns:
        The stdout from the executed code.

    Raises:
        CodeExecutionError: If the code execution fails, times out, or returns a non-zero exit code.
    """
    try:
        # Note: Using sys.executable ensures we use the same Python interpreter
        # that the main application is using.
        process = subprocess.run(
            [sys.executable, '-c', code],
            input=input_json,
            capture_output=True,
            text=True,
            timeout=timeout,
            check=True  # This will raise CalledProcessError on non-zero exit codes
        )
        logger.info("Code execution successful.")
        return process.stdout
    except subprocess.TimeoutExpired as e:
        error_message = f"Code execution timed out after {timeout} seconds."
        stderr = e.stderr if e.stderr else "No stderr output."
        logger.error(f"{error_message}\nStderr: {stderr}")
        raise CodeExecutionError(error_message, stderr) from e
    except subprocess.CalledProcessError as e:
        error_message = f"Code execution failed with exit code {e.returncode}."
        stderr = e.stderr if e.stderr else "No stderr output."
        logger.error(f"{error_message}\nStderr: {stderr}")
        raise CodeExecutionError(error_message, stderr) from e
    except Exception as e:
        error_message = f"An unexpected error occurred during code execution: {str(e)}"
        logger.error(error_message, exc_info=True)
        raise CodeExecutionError(error_message, "N/A") from e

